import os
from tqdm import tqdm
import cv2

from .build import DatasetLoader


@DatasetLoader.register()
class VOCDataset:
    @classmethod
    def build(cls, dataset_config):
        dataset_root = dataset_config['dataset_root']
        image_set_name = dataset_config['image_set_name']

        dataset = []

        image_set_path = os.path.join(dataset_root, 'imageSets', image_set_name+'.txt')
        annotations_dir = os.path.join(dataset_root, 'annotations')

        with open(image_set_path, 'r') as f:
            image_id_list = [line.strip() for line in f.readlines()]

        for image_id in tqdm(image_id_list):
            annotations_full_path = os.path.join(annotations_dir, image_id + '.txt')

            image_id_split = image_id.strip().split('/')
            image_full_path = os.path.join(dataset_root, 'images', image_id_split[0], image_id_split[1], 'lwir', image_id_split[2])
            img = None
            for suffix in ['.png', '.jpg']:
                img = cv2.imread(image_full_path + suffix)
                if img is not None:
                    image_full_path = image_full_path + suffix
                    break
            if img is None:
                raise FileNotFoundError("Can not load image " + image_full_path)

            height, width, channels = img.shape

            with open(annotations_full_path, 'r', encoding='utf-8') as f:
                object_annotation_list = f.readlines()[1:]

            objects_list = []
            for object_annotation in object_annotation_list:
                line = object_annotation.split()
                label = line[0]
                bounding_box = [float(item.strip()) for item in line[1:5]]
                bounding_box[2], bounding_box[3] = bounding_box[0] + bounding_box[2], bounding_box[1] + bounding_box[3]
                bounding_box = [max(bounding_box[0], 0.0), max(bounding_box[1], 0.0), min(bounding_box[2], width - 1.0), min(bounding_box[3], height - 1.0)]
                object_dict = {
                    'label': label,
                    'bounding_box': bounding_box
                }
                objects_list.append(object_dict)

            # to uniform interface
            objects = []
            for object in objects_list:
                object_dict = {
                    'box': object['bounding_box'],
                    'category': object['label']
                }
                objects.append(object_dict)

            record = {
                'path': image_full_path,
                'objects': objects,
                'params': {
                    'height': height,
                    'width': width,
                },
                'transforms': [],
            }
            dataset.append(record)

        return dataset
